import numpy as np
from scipy import stats
import matplotlib.pyplot as plt


def liner_fit(input_data, target_data):
    loop_max = 1000
    #0.001
    epsilon = 1e-3
    x_len = len(input_data)
    np.random.seed(0)
    w = np.random.randn(2)

    alpha = 0.001
    err = np.zeros(2)
    count = 0
    finish = 0

    while count < loop_max:
        count += 1
        sum_x = np.zeros(2)
        for i in range(x_len):
            dif = (np.dot(w, input_data[i]) - target_data[i]) * input_data[i]
            sum_x += dif
        w = w - alpha * sum_x
        if np.linalg.norm(w - err) < epsilon:
            finish = 1
            break
        else:
            err = w
    print('loop count: %d' % count, '\t w : %f, %f' % (w[0], w[1]))

    # slope, intercept, r_value, p_value, slope_std_error = stats.linregress(input_data, target_data)
    # print('intercept = %s slope = %s' %(intercept, slope))
    return w


if __name__ == '__main__':
    x = np.arange(0., 10., .2)
    # print(x)
    x_len = len(x)
    x0 = np.full(x_len, 1.0)
    # print(x0)
    input_data = np.vstack([x0, x]).T
    # print(input_data)
    target_data = 2 * x + 5 + np.random.randn(x_len)
    # print(target_data)
    w = liner_fit(input_data, target_data)
    plt.plot(x, target_data, 'k+')
    plt.plot(x, w[1]*x + w[0], 'r')
    plt.show()